#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>

#include <moveit/planning_scene_monitor/planning_scene_monitor.h>
#include <moveit/planning_scene/planning_scene.h>
#include <moveit/collision_detection/collision_common.h>
#include <moveit/collision_detection/world.h>
#include <moveit/robot_state/robot_state.h>
#include <Eigen/Core>

#include <string>
#include <chrono>
#include <regex>
#include <cstdio>

namespace robot_state = moveit::core;
using namespace std::chrono_literals;

static bool nameLike(const std::string& body, const std::string& key)
{
  if (body == key) return true;
  // 常见命名：link/collision_X 或 object_id/shape_X
  if (body.rfind(key + "/", 0) == 0) return true;                // 以 "key/" 开头
  if (body.size() >= key.size() &&
      body.compare(body.size()-key.size(), key.size(), key) == 0) // 以 key 结尾
    return true;
  if (body.find(key) != std::string::npos) return true;           // 包含
  return false;
}

class ContactDepthForceNode : public rclcpp::Node
{
public:
  ContactDepthForceNode() : Node("contact_depth_force_node")
  {
    end_link_      = this->declare_parameter<std::string>("end_link", "end");
    surface_link_  = this->declare_parameter<std::string>("surface_link", "wood_board");
    k_contact_     = this->declare_parameter<double>("k_contact", 30);        // N/m
    rate_hz_       = this->declare_parameter<double>("rate_hz", 50.0);
    frame_id_      = this->declare_parameter<std::string>("wrench_frame", "base_link");
    debug_pairs_   = this->declare_parameter<bool>("debug_pairs", true);
    force_acm_off_ = this->declare_parameter<bool>("force_disallow_acm", true);

    // 改为发布 WrenchStamped 到 /sim_force
    sim_wrench_pub_ = this->create_publisher<geometry_msgs::msg::WrenchStamped>("/sim_force", 10);
    // 保留原有的法向力话题
    wrench_pub_     = this->create_publisher<geometry_msgs::msg::WrenchStamped>("/contact_normal_wrench", 10);

    auto period = std::chrono::duration<double>(1.0 / std::max(1.0, rate_hz_));
    timer_ = this->create_wall_timer(std::chrono::duration_cast<std::chrono::milliseconds>(period),
                                     std::bind(&ContactDepthForceNode::tick, this));
  }

private:
  void lazyInitPSM()
  {
    if (inited_) return;

    try
    {
      psm_ = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(shared_from_this(), "robot_description");
    }
    catch (const std::exception& e)
    {
      RCLCPP_WARN(get_logger(), "PlanningSceneMonitor 构造失败：%s（稍后重试）", e.what());
      return;
    }

    if (!psm_ || !psm_->getPlanningScene())
    {
      RCLCPP_WARN(get_logger(), "PlanningScene 未就绪（robot_description 未加载？），稍后重试。");
      psm_.reset();
      return;
    }

    psm_->startStateMonitor();
    psm_->startSceneMonitor();
    psm_->startWorldGeometryMonitor();
    psm_->providePlanningSceneService();

    // —— 启动后做一次环境自检
    {
      planning_scene_monitor::LockedPlanningSceneRO scene(psm_);
      const auto& ps = scene;
      const auto& rm = ps->getRobotModel();
      bool has_end_in_robot   = rm->hasLinkModel(end_link_);
      bool has_surf_in_robot  = rm->hasLinkModel(surface_link_);
      bool has_surf_in_world  = ps->getWorld()->hasObject(surface_link_);

      RCLCPP_INFO(get_logger(),
        "自检：end[%s robot_link], surface[%s %s]",
        has_end_in_robot ? "是" : "否",
        surface_link_.c_str(),
        has_surf_in_robot ? "robot_link" : (has_surf_in_world ? "world_object" : "不存在"));

      if (!has_surf_in_robot && !has_surf_in_world)
      {
        RCLCPP_WARN(get_logger(),
          "⚠ 没在 PlanningScene 里找到名为 '%s' 的曲面（既不是机器人 link 也不是环境物体）。\n"
          "   - 若它是 URDF 里的 link，请确认该 link 含 <collision> 几何；只有 <visual> 是不行的。\n"
          "   - 若它是环境物体，请通过 /collision_object 或 PlanningScene 添加 id='%s' 的物体。",
          surface_link_.c_str(), surface_link_.c_str());
      }
    }

    // —— 强制把 ACM 中 end↔surface 设为“不允许碰撞”，防止被 ACM 忽略
    if (force_acm_off_)
    {
      planning_scene_monitor::LockedPlanningSceneRW scene(psm_);
      auto& acm = scene->getAllowedCollisionMatrixNonConst();
      acm.setEntry(end_link_, surface_link_, false);
      RCLCPP_INFO(get_logger(), "已将 ACM 中 '%s' <-> '%s' 设置为不允许碰撞（force_disallow_acm=true）。",
                  end_link_.c_str(), surface_link_.c_str());
    }

    inited_ = true;
    RCLCPP_INFO(get_logger(),
                "Depth->Force 估计器已就绪：end='%s', surface='%s', k=%.1f N/m, frame='%s', rate=%.1f Hz",
                end_link_.c_str(), surface_link_.c_str(), k_contact_, frame_id_.c_str(), rate_hz_);
  }

  void tick()
  {
    if (!inited_) { lazyInitPSM(); return; }
    if (!psm_ || !psm_->getPlanningScene()) return;

    collision_detection::CollisionRequest req;
    collision_detection::CollisionResult res;
    req.contacts = true;
    req.max_contacts = 512;
    req.max_contacts_per_pair = 128;
    req.distance = false;

    planning_scene_monitor::LockedPlanningSceneRO scene(psm_);
    const auto& ps = scene;
    const robot_state::RobotState& current_state = ps->getCurrentState();
    ps->checkCollision(req, res, current_state);

    Eigen::Vector3d F_sum = Eigen::Vector3d::Zero();
    size_t matched_contacts = 0;

    if (res.collision)
    {
      for (const auto& pair_contacts : res.contacts)
      {
        const std::string& b1 = pair_contacts.first.first;
        const std::string& b2 = pair_contacts.first.second;
        const auto& contacts_vec = pair_contacts.second;

        bool match_12 = nameLike(b1, end_link_) && nameLike(b2, surface_link_);
        bool match_21 = nameLike(b1, surface_link_) && nameLike(b2, end_link_);
        bool matched = match_12 || match_21;

        if (!matched) continue;

        for (const auto& c : contacts_vec)
        {
          if (c.depth <= 0.0) continue;

          // normal 从 body2 指向 body1；统一为 “surface -> end”
          Eigen::Vector3d n = match_12 ? c.normal : -c.normal;
          double Fn = k_contact_ * c.depth;
          F_sum += -Fn * n;
          matched_contacts++;
        }
      }

      if (matched_contacts == 0)
      {
        RCLCPP_WARN(get_logger(),
          "已检测到碰撞，但没有任何 depth>0 的接触点属于 '%s' <-> '%s'（可能名字不一致或 plate 无 collision）。",
          end_link_.c_str(), surface_link_.c_str());
      }
    }

    // // 直接打印力值（RCLCPP）
    // RCLCPP_INFO(get_logger(), "当前接触力：x = %.3f, y = %.3f, z = %.3f N",
    //             F_sum.x(), F_sum.y(), F_sum.z());

    // —— 发布到 /sim_force（WrenchStamped，仿照示例格式，并使用 printf 打印）
    geometry_msgs::msg::WrenchStamped sim_msg;
    sim_msg.header.stamp = this->now();
    sim_msg.header.frame_id = frame_id_;       // 示例中为 "base_link"，此处沿用可配置参数
    sim_msg.wrench.force.x = F_sum.x();
    sim_msg.wrench.force.y = F_sum.y();
    sim_msg.wrench.force.z = F_sum.z();
    sim_msg.wrench.torque.x = 0.0;
    sim_msg.wrench.torque.y = 0.0;
    sim_msg.wrench.torque.z = 0.0;
    sim_wrench_pub_->publish(sim_msg);

    const double t = rclcpp::Time(sim_msg.header.stamp).seconds();
    std::printf("[%.6f] Fx=%.3f  Fy=%.3f  Fz=%.3f  (|F|=%.3f)\n",
                t, F_sum.x(), F_sum.y(), F_sum.z(), F_sum.norm());
    std::fflush(stdout);

    // —— 兼容发布 /contact_normal_wrench（保留原话题）
    geometry_msgs::msg::WrenchStamped w;
    w.header.stamp = sim_msg.header.stamp;
    w.header.frame_id = frame_id_;
    w.wrench.force.x = F_sum.x();
    w.wrench.force.y = F_sum.y();
    w.wrench.force.z = F_sum.z();
    w.wrench.torque.x = 0.0;
    w.wrench.torque.y = 0.0;
    w.wrench.torque.z = 0.0;
    wrench_pub_->publish(w);
  }

private:
  // ROS
  rclcpp::TimerBase::SharedPtr timer_;
  rclcpp::Publisher<geometry_msgs::msg::WrenchStamped>::SharedPtr sim_wrench_pub_;
  rclcpp::Publisher<geometry_msgs::msg::WrenchStamped>::SharedPtr wrench_pub_;

  // MoveIt
  planning_scene_monitor::PlanningSceneMonitorPtr psm_;
  bool inited_{false};

  // 参数
  std::string end_link_, surface_link_, frame_id_;
  double k_contact_, rate_hz_;
  bool debug_pairs_, force_acm_off_;
};

int main(int argc, char** argv)
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<ContactDepthForceNode>();
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}
